from __future__ import annotations
import os
from abc import ABC, abstractmethod
from collections import defaultdict
from typing import List, Iterable, Mapping, Dict, Any, Optional, IO
import mind3d.utils.eagermot_utils.inputs.bbox as bbox
from mind3d.models.blocks.track_manager import TrackManager
import mind3d.utils.eagermot_utils.io as io
import mind3d.utils.eagermot_utils.base_dataset.mot_frame as mot_frame
from mind3d.utils.eagermot_utils.inputs.detection_2d import Detection2D
from mind3d.utils.eagermot_utils.fused_instance import FusedInstance
from mind3d.utils.eagermot_utils.transform.transformation import Transformation


class MOTSequence(ABC):
    """
    base class for sequence, used to retrive frame data

    Args:
        det_source (str): detection directory
        seg_source (str): segmentation directory
        split_dir (str): split directory
        name (str): sequence name
        frame_names (str): frame names
        save_result (bool): whether save result

    Examples:
        >>> super().__init__(det_source, seg_source, split_dir, name, frame_names,save_result) # recommended to inherit this class
    """
    def __init__(self, det_source: str, seg_source: str, split_dir: str, name: str, frame_names: Iterable[str],save_result: bool):
        self.det_source = det_source
        self.seg_source = seg_source
        self.split_dir = split_dir
        self.name = name
        self.frame_names = frame_names
        self.save=save_result
        # Image size for each camera - needed for 3D->2D projections. The dict is set in dataset-specific classes
        self.img_shape_per_cam: Dict[str, Any] = {}

        # Detections 3D {frame_name: [bboxes_3d]}
        self.dets_3d_per_frame: Dict[str, List[bbox.Bbox3d]] = {}

        # Detections 2D {frame_name: {cam_name: [bboxes_3d]}}
        self.dets_2d_multicam_per_frame: Dict[str, Dict[str, List[Detection2D]]] = {}

        # need to set its Transformation object and img_shape_per_cam in subclasses
        self.mot = TrackManager(self.cameras, self.classes_to_track)

        # det_seg_source_folder_name = f'{self.det_source}_{self.seg_source}'
        self.work_split_input_dir = os.path.join(self.split_dir, 'result')
        self.tracking_res_dir = os.path.join(self.work_split_input_dir, 'tracking')

    ##########################################################
    # Evaluation
    def preprocess(self, params: Mapping[str, Any]):
        '''
        preprocess before start a sequence

        Args:
            params (dict): algorithm parameters

        Returns:
            None
        '''
        folder_identifier = 'cleaning_0'  # add an easy identifier for code-related changes
        run_info: Dict[str, Any] = defaultdict(int)
        if self.det_source!="center_point":
            mot_3d_file = io.create_writable_file_if_new(
                self.get_results_folder_name(params, folder_identifier, "3d"), self.name,self.save)
            mot_2d_from_3d_file = io.create_writable_file_if_new(
                self.get_results_folder_name(params, folder_identifier, "2d_projected_3d"), self.name,self.save)

            if mot_3d_file is None:
                print(f'Sequence {self.name} already has results. Skipped')
                print('=====================================================================================')
                return run_info,0,0,False

            run_info["mot_3d_file"] = mot_3d_file.name.split(self.name)[0]
            run_info["mot_2d_from_3d_file"] = mot_2d_from_3d_file.name.split(self.name)[0]
            self.load_ego_motion_transforms()
            return run_info, mot_3d_file, mot_2d_from_3d_file,True
        return run_info, 0, 0, True

    def get_results_folder_name(self, params: Mapping[str, Any], folder_identifier: str, suffix: str):
        '''
        return result folder name based on param

        Args:
            params (dict): algorithm parameters
            folder_identifier (str): identifier
            suffix (str): 2d or 3d

        Returns:
            folder name (str)
        '''
        folder_suffix_full=suffix
        # folder_suffix_full = f"{variant_name_from_params(params)}_{folder_identifier}_{suffix}"
        return f"{self.tracking_res_dir}_{folder_suffix_full}"

    ##########################################################
    # Lazy getters for frame-specific data

    def get_segmentations_for_frame(self, frame_name: str) -> Dict[str, List[Detection2D]]:
        '''
        Return a dict of Detection2D for each camera for the requested frame

        Args:
            frame_name (str): frame name
        
        Returns:
            segmentations (Detection2D)
        '''
        if not self.dets_2d_multicam_per_frame:
            self.dets_2d_multicam_per_frame = self.load_detections_2d()
        return self.dets_2d_multicam_per_frame.get(frame_name, defaultdict(list))

    def get_bboxes_for_frame(self, frame_name: str,param) -> List[bbox.Bbox3d]:
        '''
        Return a list of bbox.Bbox3d for the requested frame
        
        Args:
            frame_name (str): frame name
        
        Returns:
            result (Bbox3d)
        '''
        if not self.dets_3d_per_frame:
            self.dets_3d_per_frame = self.load_detections_3d()
        return self.dets_3d_per_frame.get(frame_name, [])

    ##########################################################
    # Required methods and fields that need to be overridden by subclasses
    # This sadly results in some extra code, but is the best way to ensure compile-time errors

    @abstractmethod
    def load_ego_motion_transforms(self) -> None: pass

    @abstractmethod
    def save_ego_motion_transforms_if_new(self) -> None: pass

    @abstractmethod
    def load_detections_3d(self) -> Dict[str, List[bbox.Bbox3d]]: pass

    @abstractmethod
    def load_detections_2d(self) -> Dict[str, Dict[str, List[Detection2D]]]: pass

    @abstractmethod
    def get_frame(self, frame_name: str) -> mot_frame.MOTFrame: pass

    @property
    @abstractmethod
    def transformation(self) -> Transformation: pass

    @property
    @abstractmethod
    def cameras(self) -> List[str]: pass

    @property
    @abstractmethod
    def camera_default(self) -> str: pass

    @property
    @abstractmethod
    def classes_to_track(self) -> List[int]: pass

    @abstractmethod
    def report_mot_results(self, frame_name: str, predicted_instances: Iterable[FusedInstance],
                           mot_3d_file: IO,
                           mot_2d_from_3d_only_file: Optional[IO]) -> None:
        pass

    @abstractmethod
    def save_mot_results(self, mot_3d_file: IO,
                         mot_2d_from_3d_file: Optional[IO]) -> None:
        pass
